package com.ne.alice;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class NioReceiveServer {
    private Charset charset = Charset.forName("UTF-8");

    static class Client{
        String fileName;
        long fileLength;
        long startTime;
        InetSocketAddress remoteAddress;
        FileChannel outChannel;
        long receiveLength;

        public boolean isFinished(){
            return receiveLength>=fileLength;
        }
    }

    private ByteBuffer buffer = ByteBuffer.allocate(1024);
    Map<SelectableChannel,Client> clientMap = new HashMap<>();

    public void startServer() throws IOException {
        Selector selector = Selector.open();
        ServerSocketChannel serverChannel = ServerSocketChannel.open();
        ServerSocket serverSocket = serverChannel.socket();
        serverChannel.configureBlocking(false);
        InetSocketAddress address = new InetSocketAddress(18899);
        serverSocket.bind(address);
        serverChannel.register(selector, SelectionKey.OP_ACCEPT);
        System.out.println("serverChannel is listening...");
        while (selector.select()>0){
            Iterator<SelectionKey> it = selector.selectedKeys().iterator();
            while (it.hasNext()){
                SelectionKey key = it.next();
                if(key.isAcceptable()){
                    ServerSocketChannel server = (ServerSocketChannel) key.channel();
                    SocketChannel socketChannel = server.accept();
                    if(socketChannel==null){
                        continue;
                    }
                    socketChannel.configureBlocking(false);
                    SelectionKey selectionKey = socketChannel.register(selector, SelectionKey.OP_READ);
                    Client client = new Client();
                    client.remoteAddress = (InetSocketAddress) socketChannel.getRemoteAddress();
                    clientMap.put(socketChannel,client);
                    System.out.println(socketChannel.getRemoteAddress()+"连接成功...");
                }else if(key.isReadable()){
                    processData(key);
                }
                it.remove();
            }
        }
    }

    private void processData(SelectionKey key) throws IOException {
        Client client = clientMap.get(key.channel());
        SocketChannel socketChannel = (SocketChannel) key.channel();
        int num = 0;
        try{
            buffer.clear();
            while ((num = socketChannel.read(buffer))>0){
                buffer.flip();
                if(null == client.fileName){
                    if(buffer.capacity()<4){
                        continue;
                    }
                    int fileNameLen = buffer.getInt();
                    byte[] fileNameBytes = new byte[fileNameLen];
                    buffer.get(fileNameBytes);
                    // 文件名
                    String fileName = new String(fileNameBytes, charset);
                    File directory = new File("d://");
                    if (!directory.exists()) {
                        directory.mkdir();
                    }
                    System.out.println("NIO 传输目标 dir："+directory);
                    client.fileName = fileName;
                    String fullName = directory.getAbsolutePath() + File.separatorChar + fileName;
                    System.out.println("NIO 传输目标文件："+fullName);
                    File file = new File(fullName.trim());
                    if (!file.exists()) {
                        file.createNewFile();
                    }
                    FileChannel fileChannel = new FileOutputStream(file).getChannel();
                    client.outChannel = fileChannel;
                    if (buffer.capacity() < 8) {
                        continue;
                    }
                    // 文件长度
                    long fileLength = buffer.getLong();
                    client.fileLength = fileLength;
                    client.startTime = System.currentTimeMillis();
                    System.out.println("NIO 传输开始：");
                    client.receiveLength += buffer.capacity();
                    if (buffer.capacity() > 0) {
                        // 写入文件
                        client.outChannel.write(buffer);
                    }
                    if (client.isFinished()) {
                        finished(key, client);
                    }
                    buffer.clear();
                }
                else {
                    client.receiveLength += buffer.capacity();
                    // 写入文件
                    client.outChannel.write(buffer);
                    if (client.isFinished())
                    {
                        finished(key, client);
                    }
                    buffer.clear();
                }
            }
            key.cancel();
        } catch (IOException e) {
            key.cancel();
            e.printStackTrace();
            return;
        }
        // 调用 close 为-1 到达末尾
        if (num == -1) {
            finished(key, client);
            buffer.clear();
        }
    }

    private void finished(SelectionKey key, Client client) throws IOException {
        client.outChannel.close();
        System.out.println("上传完毕");
        key.cancel();
        System.out.println("文件接收成功,File Name：" + client.fileName);
        System.out.println(" Size：" + (client.fileLength));
        long endTime = System.currentTimeMillis();
        System.out.println("NIO IO 传输毫秒数：" + (endTime - client.startTime));
    }

    public static void main(String[] args) throws IOException {
        NioReceiveServer server = new NioReceiveServer();
        server.startServer();
    }

}
